package runtime

import (
	"context"
	"errors"
	"fmt"
	"maps"
	"math"
	"reflect"
	"slices"
	"strings"

	"tvl.fyi/tvix/eval/builtins"
	"tvl.fyi/tvix/eval/internal/bytecode"
	"tvl.fyi/tvix/eval/internal/compiler"
	"tvl.fyi/tvix/eval/internal/index"
	"tvl.fyi/tvix/eval/internal/op"
	"tvl.fyi/tvix/eval/value"
)

type frame struct {
	chunk    *bytecode.Chunk
	ip       int
	offset   uint64
	captures []value.NixValue
}

type runtime struct {
	frames []frame
	stack  []value.NixValue
}

func (r *runtime) push(val value.NixValue) {
	r.stack = append(r.stack, val)
}

func (r *runtime) pop() (val value.NixValue) {
	val, r.stack = r.stack[len(r.stack)-1], r.stack[:len(r.stack)-1]
	return
}

func (r *runtime) peek() value.NixValue {
	return r.stack[len(r.stack)-1]
}

// runtimeContext returns a context that is equipped to handle nested
// evaluation.
//
// TODO: thread through outer context; make this good
func (r *runtime) context() context.Context {
	return context.WithValue(context.Background(), bytecode.RuntimeKey{}, &runtime{})
}

func typed[T value.NixValue](v value.NixValue) (T, error) {
	t, ok := v.(T)
	if !ok {
		if rv := reflect.ValueOf(t); !rv.IsValid() {
			// TODO: two-levels of indirection -> bad type error :/
			return t, typeError{v.NixType(), "another type"}
		}
		return t, typeError{v.NixType(), t.NixType()}
	}

	return t, nil
}

func (r *runtime) runFrame() error {
	var err error
	var f frame
	f, r.frames = r.frames[len(r.frames)-1], r.frames[:len(r.frames)-1]

	for f.ip < len(f.chunk.Code) {
		next := f.chunk.ReadOp(f.ip)
		switch next {
		case op.Constant:
			arg, size := f.chunk.ReadUint(f.ip + 1)
			f.ip += size
			r.push(f.chunk.Constants[arg])

		case op.Pop:
			r.pop()

		case op.Force:
			t, isThunk := r.peek().(*bytecode.Thunk)
			for isThunk {
				r.stack[len(r.stack)-1], err = t.Force(r.context())
				if err != nil {
					return err
				}
				t, isThunk = r.peek().(*bytecode.Thunk)
			}

		case op.Invert:
			b, err := typed[value.Bool](r.pop())
			if err != nil {
				return err
			}

			r.push(!b)

		case op.Negate:
			n, err := typed[value.Number](r.pop())
			if err != nil {
				return err
			}

			r.push(n.Mul(value.Int(-1)))

		case op.Add:
			rhs := r.pop()
			lhs := r.pop()
			lhsNum, lhsOk := lhs.(value.Number)
			rhsNum, rhsOk := rhs.(value.Number)
			if lhsOk && rhsOk {
				r.push(lhsNum.Add(rhsNum))
			} else {
				lhsStr, lhsOk := lhs.(value.String)
				rhsStr, rhsOk := rhs.(value.String)
				if lhsOk && rhsOk {
					r.push(lhsStr + rhsStr)
				} else {
					return fmt.Errorf("cannot add %q to %q", rhs.NixType(), lhs.NixType())
				}
			}

		case op.Sub:
			if err := r.runBinop("subtract", value.Number.Sub); err != nil {
				return err
			}

		case op.Mul:
			if err := r.runBinop("multiply", value.Number.Mul); err != nil {
				return err
			}

		case op.Div:
			if err := r.runBinop("divide", value.Number.Div); err != nil {
				return err
			}

		case op.Equal:
			rhs := r.pop()
			lhs := r.pop()

			// TODO: until proper comparing is in place, check for comparability
			// first.
			if !(reflect.ValueOf(rhs).Comparable() && reflect.ValueOf(lhs).Comparable()) {
				r.push(value.Bool(false))
			} else {
				r.push(value.Bool(lhs == rhs))
			}

		case op.List:
			arg, size := f.chunk.ReadUint(f.ip + 1)
			f.ip += size

			if arg == 0 {
				r.push(value.NixList([]value.NixValue{}))
				break
			}

			start := len(r.stack) - int(arg)
			elems := make([]value.NixValue, arg)
			copy(elems, r.stack[start:])
			r.stack = r.stack[:start]
			r.push(value.NixList(elems))

		case op.Concat:
			rhs, err := typed[value.NixList](r.pop())
			if err != nil {
				return err
			}

			lhs, err := typed[value.NixList](r.pop())
			if err != nil {
				return err
			}

			r.push(slices.Concat(lhs, rhs))

		case op.Update:
			rhs, err := typed[value.NixAttrs](r.pop())
			if err != nil {
				return err
			}

			lhs, err := typed[value.NixAttrs](r.pop())
			if err != nil {
				return err
			}

			merged := make(value.NixAttrs)
			maps.Copy(merged, rhs)
			maps.Copy(merged, lhs)
			r.push(merged)

		case op.Interpolate:
			arg, size := f.chunk.ReadUint(f.ip + 1)
			f.ip += size

			start := len(r.stack) - int(arg)
			var buf strings.Builder
			for _, fragment := range r.stack[start:] {
				s, err := typed[value.String](fragment)
				if err != nil {
					return err
				}

				buf.WriteString(string(s))
			}

			r.stack = r.stack[:start]
			r.push(value.String(buf.String()))

		case op.Attrs:
			arg, size := f.chunk.ReadUint(f.ip + 1)
			f.ip += size

			attrs := make(map[value.String]value.NixValue, arg)
			for range arg {
				val := r.pop()
				key, err := typed[value.String](r.pop())
				if err != nil {
					return err
				}

				attrs[key] = val
			}

			r.push(value.NixAttrs(attrs))

		case op.Select:
			key, err := typed[value.String](r.pop())
			if err != nil {
				return err
			}

			attrs, err := typed[value.NixAttrs](r.pop())
			if err != nil {
				return err
			}

			val, found := attrs[key]
			if !found {
				return fmt.Errorf("key %s not found in attribute set", key)
			}

			r.push(val)

		case op.CloseScope:
			arg, size := f.chunk.ReadUint(f.ip + 1)
			f.ip += size

			// Immediately move the top value into the right position.
			targetIdx := len(r.stack) - 1 - int(arg)
			r.stack[targetIdx] = r.pop()

			// Then drop the remaining values.
			r.stack = r.stack[:targetIdx+1]

		case op.GetLocal:
			arg, size := f.chunk.ReadUint(f.ip + 1)
			f.ip += size
			r.push(r.stack[f.offset+arg])

		case op.GetCapture:
			/* captures need to wired up correctly first
			arg, size := f.chunk.ReadUint(f.ip + 1)
			f.ip += size
			r.push(f.captures[arg]) */
			return fmt.Errorf("capture access not fully supported yet")

		case op.Call:
			switch target := r.pop().(type) {
			case compiler.Callable:
				f.ip++

				// Push the current frame back, then create & enter the call
				// frame.
				//
				// TODO: with thunked applications this might be unnecessary.
				r.frames = append(r.frames, f, frame{
					chunk:  target.Chunk(),
					offset: uint64(len(r.stack) - 1),
				})
				return nil
			case builtins.NixBuiltin:
				val, err := target.Call(context.TODO(), r.pop())
				if err != nil {
					return fmt.Errorf("failed to call builtin %q: %w", target.Name(), err)
				}

				r.push(val)

			default:
				return fmt.Errorf("expected a callable value, but got %s", target.NixType())
			}

		case op.Jump:
			f.ip += int(f.chunk.ReadJumpOffset(f.ip))
			continue

		case op.JumpIfFalse:
			top, err := typed[value.Bool](r.peek())
			if err != nil {
				return err
			}

			if top {
				f.ip += 2 // offset size
			} else {
				f.ip += int(f.chunk.ReadJumpOffset(f.ip))
				continue
			}

		case op.JumpIfTrue:
			top, err := typed[value.Bool](r.peek())
			if err != nil {
				return err
			}

			if !top {
				f.ip += 2 // offset size
			} else {
				f.ip += int(f.chunk.ReadJumpOffset(f.ip))
				continue
			}

		case op.AssertBool:
			_, ok := r.peek().(value.Bool)
			if !ok {
				return typeError{r.peek().NixType(), "bool"}
			}

		case op.Capture:
			count, size := f.chunk.ReadUint(f.ip + 1)
			f.ip += size

			// first actual argument is the stack index
			idx, size := f.chunk.ReadUint(f.ip + 1)
			f.ip += size

			t, isThunk := r.stack[idx].(*bytecode.Thunk)
			if !isThunk {
				return errors.New("tvix bug: capture target is not a thunk")
			}

			captures := make([]value.NixValue, count-1)
			for n := range count - 1 {
				// TODO: how to actually capture?
				// stack index: easy, just take it from the stack
				// capture index: ??, where are the current captures?
				posInt, size := f.chunk.ReadUint(f.ip + 1)
				f.ip += size
				pos := index.Position(posInt)

				if stackPos, ok := pos.RuntimeStackIndex(); ok {
					captures[n] = r.stack[index.Stack(f.offset)+stackPos]
					continue
				}

				if _, ok := pos.CaptureIndex(); ok {
					return errors.New("tvix bug: capturing captures not yet implemented")
				}

				return errors.New("tvix bug: invalid capture position encountered")
			}

			err := t.Capture(captures)
			if err != nil {
				return err
			}

		case op.Return:
			return nil

		default:
			return fmt.Errorf("unsupported op: %s", next)
		}

		f.ip++
	}

	return nil
}

var (
	divByZero = errors.New("can not divide by zero")
	infinity  = value.Float(math.Inf(1))
)

func (r *runtime) runBinop(verb string, op func(value.Number, value.Number) value.Number) (err error) {
	// TODO: solve this some other way
	defer func() {
		if p := recover(); p != nil {
			err = divByZero
		}
	}()

	rhs := r.pop()
	lhs := r.pop()
	lhsNum, lhsOk := lhs.(value.Number)
	rhsNum, rhsOk := rhs.(value.Number)
	if lhsOk && rhsOk {
		result := op(lhsNum, rhsNum)
		if result == infinity {
			err = divByZero
		}
		r.push(result)
		return
	}

	err = fmt.Errorf("cannot %s %q and %q", verb, lhs.NixType(), rhs.NixType())
	return
}

func (r *runtime) run() error {
	for len(r.frames) > 0 {
		err := r.runFrame()
		if err != nil {
			return err
		}
	}

	return nil
}

func (r *runtime) Evaluate(c *bytecode.Chunk, captures []value.NixValue) (value.NixValue, error) {
	r.frames = append(r.frames[:0], frame{
		chunk:    c,
		captures: captures,
	})
	r.stack = r.stack[:0]

	err := r.run()
	if err != nil {
		return nil, err
	}

	return r.pop(), nil
}

func RunChunk(c *bytecode.Chunk, strict bool) (value.NixValue, error) {
	r := runtime{}
	val, err := r.Evaluate(c, nil)
	if strict && err == nil {
		err = val.Strict(r.context())
	}

	return val, err
}
